from fastNLP import BucketedBatchSampler
import numpy as np


class GradualBucketedBatchSampler(BucketedBatchSampler):
    def __init__(self, dataset, length, batch_size: int = 32, num_batch_per_bucket: int = 10,
                 shuffle: bool = True, drop_last: bool = False, seed: int = 0, **kwargs):
        super(GradualBucketedBatchSampler, self).__init__(dataset, length, batch_size=batch_size,
                                                          num_batch_per_bucket=num_batch_per_bucket,
                                                          shuffle=shuffle, drop_last=drop_last, seed=seed, **kwargs)

    def bucketerize(self, sorted_indices, batch_size, num_batch_per_bucket, seed):
        """
        将 ``indices`` 分桶

        :param sorted_indices: List[int]
        :param batch_size: int
        :param num_batch_per_bucket: int
        :param seed: int
        :return:
        """
        # 实际的 bucket 大小
        bucket_size = min(len(sorted_indices), batch_size * num_batch_per_bucket)
        rng = np.random.default_rng(abs(seed))
        num_buckets = (len(sorted_indices) + bucket_size - 1) // bucket_size
        batches = []
        batch_indices = []
        for i in range(num_buckets):
            bucket = sorted_indices[i * bucket_size:(i + 1) * bucket_size]
            rng.shuffle(bucket)  # bucket 内部 shuffle 一下
            _num_batches = len(bucket) // batch_size
            if _num_batches == 0:
                _batches = [bucket]
            else:
                _batches = np.array_split(bucket[:_num_batches * batch_size], _num_batches)
                if len(bucket) % batch_size != 0:
                    _batches.append(bucket[_num_batches * batch_size:])
            batch_indices.extend(list(range(len(batches), len(batches) + len(_batches))))
            batches.extend(_batches)
        last_batches = []
        # 最后一个batch 统一不参与shuffle，因为有的rank最后一个 batch 可能不足一个batch_size （不足的时候
        #  一定要放在末尾，所以就干脆所有的rank都不对最后一个batch进行shuffle）。
        if len(batches) >= 1:
            last_batches = [list(batches[-1])]
        batch_indices = list(batch_indices[:-1])
        rng = np.random.default_rng(abs(seed))  # 这里防止由于bucket长度不同，对随机数状态有影响
        rng.shuffle(batch_indices)  # 不同的 batch 也 shuffle ，当前这种可以保证每张卡上每个 batch 长度都接近的。
        batches = (np.array(batches, dtype=object)[batch_indices]).tolist()
        if last_batches:
            batches = batches + last_batches
        batches.sort(lambda x: sum(x))
        return batches
